import segyio
import matplotlib.pyplot as plt
import numpy as np


def plot_corridors(file1, file2):
    try:
        # 读取第一个 SGY 文件
        with segyio.open(file1, 'r', ignore_geometry=True) as segy1:
            data1 = segy1.trace.raw[:]

        # 读取第二个 SGY 文件
        with segyio.open(file2, 'r', ignore_geometry=True) as segy2:
            data2 = segy2.trace.raw[:]

        # 获取第一个走廊的道数
        num_traces1 = data1.shape[0]
        # 获取第二个走廊的道数
        num_traces2 = data2.shape[0]

        # 调整第一个走廊的位置，使其最后一道与第二个走廊的第 13 道相邻
        offset = 13 - num_traces1

        # 确保 offset 有效
        if offset < 0:
            offset = 0

        fig = plt.figure()
        # 创建图形和两个子图（上下排列）
        ax1 = fig.add_axes([0.1, 0.2, 0.8, 0.5])  # main axes
        ax2 = fig.add_axes([0.1, 0.1, 0.8, 0.5])  # inset axes
        ax2.set_facecolor((1, 1, 1, 0))  # RGBA 颜色，其中 A（透明度）为 0 表示完全透明

        # 绘制第一个走廊的数据
        for trace_num in range(num_traces1):
            trace_data = data1[trace_num, :]
            # 过滤掉值为 0 的数据
            non_zero_indices = np.nonzero(trace_data)[0]
            non_zero_data = trace_data[non_zero_indices]
            # 绘制非零数据，设置透明度为 1（不透明）
            ax1.plot(non_zero_indices, non_zero_data + trace_num, color='black', alpha=1)

        # 绘制第二个走廊的数据
        for trace_num in range(num_traces2):
            trace_data = data2[trace_num, :]
            # 过滤掉值为 0 的数据
            non_zero_indices = np.nonzero(trace_data)[0]
            non_zero_data = trace_data[non_zero_indices]
            # 绘制非零数据，设置透明度为 1（不透明）
            ax2.plot(non_zero_indices, non_zero_data + trace_num, color='black', alpha=1)

        # 设置坐标轴标签和标题
        ax1.set_ylabel('Trace Number')
        ax1.set_title('Corridor 1')

        ax2.set_xlabel('Time Sample')
        ax2.set_ylabel('Trace Number')
        ax2.set_title('Corridor 2')

        # 调整子图之间的间距，使两个走廊靠近
        plt.subplots_adjust(hspace=0.05)

        plt.show()

    except Exception as e:
        print(f"An error occurred: {e}")



def plot():
    import matplotlib.pyplot as plt
    import numpy as np
    import math

    x = np.arange(0, math.pi * 2, 0.05)
    y = np.sin(x)

    fig = plt.figure()
    axes1 = fig.add_axes([0.1, 0.1, 0.8, 0.8])  # main axes
    axes2 = fig.add_axes([0.55, 0.55, 0.3, 0.3])  # inset axes
    axes1.plot(x, y, 'b')
    axes2.plot(x, np.cos(x), 'r')
    axes1.set_title('sine')
    axes2.set_title("cosine")

    plt.show()


# 使用示例
stkb_path = r"E:\vspData\斜井\stkb.sgy"  # 滤波走廊
st_path = r"E:\vspData\斜井\st.sgy"  # 走廊
plot_corridors(stkb_path, st_path)
